/*****************************************************************************
 * load.c : load and parse coff and library file
 *****************************************************************************
 * Copyright (C) 2011 PELinker
 *
 * Authors: Iceberg <iceberg58744@163.com>
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston MA 02110-1301, USA.
 *****************************************************************************/

#include "common.h"
#include "memory.h"
#include "util.h"

#define LEADING_CHARARECTER '/'

#define PADDING_CHARARECTER ' '

#define LIB_SIGNATURE "!<arch>\n"

static int ConvertDecimalStringToDWORD(const BYTE * const src, DWORD *dest)
{
	DWORD i = 0;

	*dest = 0;

	while (src[i])
	{
		if (src[i] >= '0' && src[i] <= '9')
		{
			*dest *= 10;
			*dest += (src[i] - '0');
		}
		else
		{
			return ERR_INVALID_ARGUMENT;
		}

		++i;
	}

	return 0;
}

static int LoadCOFF(const HANDLE hFile, Link_Data * const link_data)
{
	HANDLE hFileMapping;
	LPVOID pFileData;
	const BYTE *FileContent;
	const IMAGE_FILE_HEADER *pCOFFHeader;
	const IMAGE_SECTION_HEADER *pSectionHeader;
	const IMAGE_SYMBOL *pSymbolTable;
	const BYTE *pStringTable;
	DWORD StringTableSize;
	Section *pSection;
	Symbol *pSymbol;
	DWORD i;
	DWORD j;

	hFileMapping = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
	if (NULL == hFileMapping)
	{
		return ERR_INVALID_FILEHANDLE;
	}

	pFileData = MapViewOfFile(hFileMapping, FILE_MAP_READ, 0, 0, 0);
	if (NULL == pFileData)
	{
		CloseHandle(hFileMapping);

		return ERR_INVALID_FILEHANDLE;
	}

	FileContent = (const BYTE *)pFileData;

	pCOFFHeader = (IMAGE_FILE_HEADER *)FileContent;

	pSectionHeader = (IMAGE_SECTION_HEADER *)(FileContent + sizeof(IMAGE_FILE_HEADER));

	pSymbolTable = (IMAGE_SYMBOL *)(FileContent + pCOFFHeader->PointerToSymbolTable);

	pStringTable = FileContent + pCOFFHeader->PointerToSymbolTable + pCOFFHeader->NumberOfSymbols * sizeof(IMAGE_SYMBOL);

	StringTableSize = *((DWORD *)pStringTable);

	link_data->objects[link_data->NumberOfObjects].NumberOfSections = pCOFFHeader->NumberOfSections;
	link_data->objects[link_data->NumberOfObjects].NumberOfSymbols = pCOFFHeader->NumberOfSymbols;
	link_data->objects[link_data->NumberOfObjects].sections = (Section *)Malloc(sizeof(Section) * pCOFFHeader->NumberOfSections);
	link_data->objects[link_data->NumberOfObjects].symbols = (Symbol *)Malloc(sizeof(Symbol) * pCOFFHeader->NumberOfSymbols);

	pSection = link_data->objects[link_data->NumberOfObjects].sections;

	pSymbol = link_data->objects[link_data->NumberOfObjects].symbols;

	for (i = 0; i < pCOFFHeader->NumberOfSections; ++i)
	{
		if (LEADING_CHARARECTER == pSectionHeader[i].Name[0])
		{
			DWORD index;
			DWORD len;

			if (ConvertDecimalStringToDWORD(pSectionHeader[i].Name + 1, &index) || index > StringTableSize)
			{
				UnmapViewOfFile(pFileData);
				CloseHandle(hFileMapping);

				return ERR_INVALID_FILETYPE;
			}

			len = lstrlen(pStringTable + index);

			pSection[i].Name = (BYTE *)Malloc(len + 1);
			Memcpy(pSection[i].Name, pStringTable + index, len);
		}
		else
		{
			pSection[i].Name = (BYTE *)Malloc(IMAGE_SIZEOF_SHORT_NAME + 1);
			Memcpy(pSection[i].Name, pSectionHeader[i].Name, IMAGE_SIZEOF_SHORT_NAME);
		}

		if (pSectionHeader[i].SizeOfRawData && pSectionHeader[i].PointerToRawData)
		{
			pSection[i].PointerToRawData = (BYTE *)Malloc(pSectionHeader[i].SizeOfRawData);
			Memcpy(pSection[i].PointerToRawData, FileContent + pSectionHeader[i].PointerToRawData, pSectionHeader[i].SizeOfRawData);
		}

		pSection[i].SizeOfRawData = pSectionHeader[i].SizeOfRawData;

		if (pSectionHeader[i].PointerToRelocations && pSectionHeader[i].NumberOfRelocations)
		{
			pSection[i].PointerToRelocations = (Relocation *)Malloc(pSectionHeader[i].NumberOfRelocations * sizeof(Relocation));

			for (j = 0; j < pSectionHeader[i].NumberOfRelocations; ++j)
			{
				pSection[i].PointerToRelocations[j].VirtualAddress = ((IMAGE_RELOCATION *)(FileContent + pSectionHeader[i].PointerToRelocations))[j].VirtualAddress;
				pSection[i].PointerToRelocations[j].SymbolTableIndex = ((IMAGE_RELOCATION *)(FileContent + pSectionHeader[i].PointerToRelocations))[j].SymbolTableIndex;
				pSection[i].PointerToRelocations[j].Type = ((IMAGE_RELOCATION *)(FileContent + pSectionHeader[i].PointerToRelocations))[j].Type;
				pSection[i].PointerToRelocations[j].PointerToSection = pSection + i;
				pSection[i].PointerToRelocations[j].PointerToSymbolTable = pSymbol;
			}

			pSection[i].NumberOfRelocations = pSectionHeader[i].NumberOfRelocations;
		}

		pSection[i].Characteristics = pSectionHeader[i].Characteristics;
	}

	for (i = 0; i < pCOFFHeader->NumberOfSymbols; ++i)
	{
		if (pSymbolTable[i].N.Name.Short)
		{
			pSymbol[i].Name = (BYTE *)Malloc(IMAGE_SIZEOF_SHORT_NAME + 1);
			Memcpy(pSymbol[i].Name, pSymbolTable[i].N.ShortName, IMAGE_SIZEOF_SHORT_NAME);
		}
		else
		{
			DWORD len;

			if (pSymbolTable[i].N.Name.Long > StringTableSize)
			{
				UnmapViewOfFile(pFileData);
				CloseHandle(hFileMapping);

				return ERR_INVALID_FILETYPE;
			}

			len = lstrlen(pStringTable + pSymbolTable[i].N.Name.Long);

			pSymbol[i].Name = (BYTE *)Malloc(len + 1);
			Memcpy(pSymbol[i].Name, pStringTable + pSymbolTable[i].N.Name.Long, len);
		}

		pSymbol[i].Value = pSymbolTable[i].Value;

		if (pSymbolTable[i].SectionNumber > 0)
		{
			pSymbol[i].PointerToSection = pSection + pSymbolTable[i].SectionNumber - 1;
			pSymbol[i].Offset = pSymbolTable[i].Value;
		}

		pSymbol[i].SectionNumber = pSymbolTable[i].SectionNumber;
		pSymbol[i].StorageClass = pSymbolTable[i].StorageClass;
		pSymbol[i].NumberOfAuxSymbols = pSymbolTable[i].NumberOfAuxSymbols;

		i += (pSymbol[i].NumberOfAuxSymbols);
	}

	++(link_data->NumberOfObjects);

	UnmapViewOfFile(pFileData);
	CloseHandle(hFileMapping);

	return 0;
}

static int LoadLib(const HANDLE hFile, Link_Data * const link_data)
{
	HANDLE hFileMapping;
	LPVOID pFileData;
	const BYTE *FileContent;
	const BYTE *SecondLinkerMember;
	const IMAGE_ARCHIVE_MEMBER_HEADER *pLibraryHeader;
	BYTE *pStringTable;
	DWORD Size;
	DWORD HeaderSize;
	DWORD NumberOfMembers;
	DWORD NumberOfSymbols;
	DWORD i;
	DWORD j;

	hFileMapping = CreateFileMapping(hFile, NULL, PAGE_READONLY, 0, 0, NULL);
	if (NULL == hFileMapping)
	{
		return ERR_INVALID_FILEHANDLE;
	}

	pFileData = MapViewOfFile(hFileMapping, FILE_MAP_READ, 0, 0, 0);
	if (NULL == pFileData)
	{
		CloseHandle(hFileMapping);

		return ERR_INVALID_FILEHANDLE;
	}

	FileContent = (const BYTE *)pFileData;

	pLibraryHeader = (const IMAGE_ARCHIVE_MEMBER_HEADER *)(FileContent + IMAGE_ARCHIVE_START_SIZE);

	Size = 0;
	i = 0;
	while (pLibraryHeader->Size[i] != PADDING_CHARARECTER)
	{
		Size *= 10;
		Size += (pLibraryHeader->Size[i] - '0');
		++i;
	}

	pLibraryHeader = (const IMAGE_ARCHIVE_MEMBER_HEADER *)(FileContent + IMAGE_ARCHIVE_START_SIZE + IMAGE_SIZEOF_ARCHIVE_MEMBER_HDR + Size);

	SecondLinkerMember = FileContent + IMAGE_ARCHIVE_START_SIZE + IMAGE_SIZEOF_ARCHIVE_MEMBER_HDR + Size + IMAGE_SIZEOF_ARCHIVE_MEMBER_HDR;

	Size = 0;
	i = 0;
	while (pLibraryHeader->Size[i] != PADDING_CHARARECTER)
	{
		Size *= 10;
		Size += (pLibraryHeader->Size[i] - '0');
		++i;
	}

	NumberOfMembers = *((const DWORD *)(SecondLinkerMember));

	NumberOfSymbols = *((const DWORD *)(SecondLinkerMember + sizeof(NumberOfMembers) + 4 * NumberOfMembers));

	link_data->libraries[link_data->NumberOfLibraries].NumberOfSymbols = NumberOfSymbols;

	link_data->libraries[link_data->NumberOfLibraries].SymbolNames = (BYTE **)Malloc(sizeof(BYTE *) * NumberOfSymbols);

	link_data->libraries[link_data->NumberOfLibraries].ReferenceNames = (BYTE **)Malloc(sizeof(BYTE *) * NumberOfSymbols);

	HeaderSize = sizeof(NumberOfMembers) + 4 * NumberOfMembers + sizeof(NumberOfSymbols) + 2 * NumberOfSymbols;

	pStringTable = (BYTE *)Malloc(Size - HeaderSize);
	Memcpy(pStringTable, SecondLinkerMember + HeaderSize, Size - HeaderSize);

	j = 0;
	for (i = 0; i < NumberOfSymbols; ++i)
	{
		link_data->libraries[link_data->NumberOfLibraries].SymbolNames[i] = pStringTable + j;
		while (pStringTable[j])
		{
			++j;
		}
		++j;
	}

	i = 0;
	while ((SecondLinkerMember + HeaderSize)[j] != LEADING_CHARARECTER)
	{
		link_data->libraries[link_data->NumberOfLibraries].LibraryName[i] = (SecondLinkerMember + HeaderSize)[j];
		++i;
		++j;
	}
	link_data->libraries[link_data->NumberOfLibraries].LibraryName[i] = '\0';

	++(link_data->NumberOfLibraries);

	UnmapViewOfFile(pFileData);
	CloseHandle(hFileMapping);

	return 0;
}

static int LoadFile(const String * const FileName, Link_Data * const link_data)
{
	int ret = ERR_INVALID_FILETYPE;
	char *name = (char *)Malloc(FileName->length + 1);
	IMAGE_FILE_HEADER coff;
	HANDLE hFile;
	DWORD FileSize;

	Memcpy(name, FileName->ptr, FileName->length);
	name[FileName->length] = '\0';

	hFile = CreateFile(name, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);

	if (INVALID_HANDLE_VALUE == hFile)
	{
		PrintInfoAux("Cannot open file: ");
		PrintInfo(name);
		PrintInfoAux("\n");

		return ERR_INVALID_FILENAME;
	}

	FileSize = GetFileSize(hFile, NULL);

	if (FileSize < sizeof(coff))
	{
		PrintInfoAux("File is incomplete: ");
		PrintInfo(name);
		PrintInfoAux("\n");
		goto err;
	}

	if (!ReadFile(hFile, &coff, sizeof(coff), &FileSize, NULL))
	{
		PrintInfoAux("Cannot read file: ");
		PrintInfo(name);
		PrintInfoAux("\n");
		goto err;
	}

	if (IMAGE_FILE_MACHINE_I386 == coff.Machine && 0 == coff.Characteristics)
	{
		if (coff.SizeOfOptionalHeader)
		{
			PrintInfoAux("Cannot process linked file: ");
			PrintInfo(name);
			PrintInfoAux("\n");
			goto err;
		}

		if (LoadCOFF(hFile, link_data))
		{
			PrintInfoAux("Load obj file failed: ");
			PrintInfo(name);
			PrintInfoAux("\n");
			goto err;
		}
	}
	else if (!Memcmp(&coff, LIB_SIGNATURE, sizeof(LIB_SIGNATURE) - 1))
	{
		if (LoadLib(hFile, link_data))
		{
			PrintInfoAux("Load lib file failed: ");
			PrintInfo(name);
			PrintInfoAux("\n");
			goto err;
		}
	}
	else
	{
		PrintInfoAux("Unsupported file type: ");
		PrintInfo(name);
		PrintInfoAux("\n");
		goto err;
	}

	ret = 0;
err:
	CloseHandle(hFile);

	return ret;
}

int LoadFiles(const Command * const cmd, Link_Data * const link_data)
{
	DWORD i;

	if (!cmd->TotalFileNumber)
	{
		return ERR_MISSING_FILENAME;
	}

	for (i = 0; i < cmd->TotalFileNumber; ++i)
	{
		int ret;
#ifdef CHECK_DUPLICATE_FILE
		DWORD j;

		for (j = 0; j < i; ++j)
		{
			if (!StringCompare(&(cmd->files[i]), &(cmd->files[j])))
			{
				PrintInfoAux("Duplicated file: ");
				PrintInfoEx(cmd->files[i].ptr, cmd->files[i].length);
				PrintInfoAux("\n");

				return ERR_DUPLICATED_FILENAME;
			}
		}
#endif
		ret = LoadFile(&(cmd->files[i]), link_data);

		if (ret)
		{
			return ret;
		}
	}

	return 0;
}
